#include "Config.h"

#ifdef _MPI_IMPLEMENT
#include "MPIPackage.h"
#include "CMASearcherMPI.h"
#include "cmaes_interface.h"
#include <stdio.h>
#include "SceneBuilder.h"
#include "Scene.h"
#include "DecoArticulatedBody.h"
#include "timer.h"
#include "DecoLogger.h"
#include "CMAContinuationMPI.h"
#include "ConfigManager.h"
#include "SearcherConfig.h"

extern double (*gEvaluator)(double*, DecoScene*, double*);

void LogPopData(FILE* fpLog, const char* log)
{
	fprintf(fpLog, log);
	fflush(fpLog);
}

#define BUFF_SIZE 128

CMASearcherMPI::CMASearcherMPI(int dim, int assigned) : CMASearcher(dim, assigned)
{
}

CMASearcherMPI::~CMASearcherMPI()
{

}
int CMASearcherMPI::Search(double* lower_bound, double* upper_bound, double* argMin, int maxIterations, int rank, int numProcs, int* pSelectedNodes)
{
	if (rank == 0)
		master(lower_bound, upper_bound, argMin, maxIterations, rank, numProcs, pSelectedNodes);
	else
		worker(rank);
	return 1;
}

void CMASearcherMPI::master(double* lower_bound, double* upper_bound, double* argMin, int maxIterations, int rank, int numProcs, int* pSelectedNodes)
{
	//	printf("Start CMASearchMPI master function at rank %d\n", rank);
	setInitialGuess(lower_bound, upper_bound);

	int iterationCount = 0;

	cmaes_t evo; /* an CMA-ES type struct or "object" */
	double *arFunvals, *const*pop, *xfinal;
	const SearcherInfo& info = DecoConfig::GetSingleton()->GetSearcherInfo();
	/* Initialize everything into the struct evo, 0 means default */
	arFunvals = cmaes_init(&evo, mDim, mPrevSol, mStandardDeviation, 0, 0, info.mCMAInitialPath.c_str()); 
	cmaes_WriteLogFile(cmaes_SayHello(&evo));
	cmaes_WriteLogFile("\n");
	cmaes_ReadSignals(&evo, info.mCMASignalPath.c_str());  /* write header and initial values */

	/* Iterate until stop criterion holds */
	int	numSamples = cmaes_Get(&evo, "lambda");

	FILE* fplog = fopen("PopData.txt", "a");
	while(!cmaes_TestForTermination(&evo))
	{ 
		/* generate lambda new search points, sample population */
		pop = cmaes_SamplePopulation(&evo); /* do not change content of pop */

		for (int i = 0; i < cmaes_Get(&evo, "popsize"); ++i) 
		{
			while (!isFeasible(pop[i], lower_bound, upper_bound, info.mJointConstraints))
				cmaes_ReSampleSingle(&evo, i);
		}
		numSamples = cmaes_Get(&evo, "lambda");
		DecoTimer timer;
		timer.startTimer();
		int numPasses = numSamples / (numProcs - 1) + ((numSamples % (numProcs - 1)) ? 1 : 0);
		char buff[BUFF_SIZE];
		char logStr[512];
		printf("NumPasses: %d, NumSamples: %d\n", numPasses, numSamples);
		for (int ithPass = 0; ithPass < numPasses; ++ithPass)
		{
			for (int i = 0; i < numProcs - 1; ++i)
			{
				int ithPop = ithPass * (numProcs - 1) + i;
				if (ithPop >= numSamples)
					break;
				int dest = pSelectedNodes[i];
				int tag = pSelectedNodes[i];
				Package pkg(mDim, pop[ithPop]);
				pkg.Pack(buff, BUFF_SIZE);

				MPI_Send(&buff, BUFF_SIZE, MPI_PACKED, dest, tag, MPI_COMM_WORLD);
			}
			printf("All messages sent.\n");

			vector<NodeTimeTuple> vTuple;
			for (int i = 0; i < numProcs - 1; ++i)
			{
				int ithPop = ithPass * (numProcs - 1) + i;
				if (ithPop >= numSamples)
					break;
				int src = pSelectedNodes[i];
				int tag = pSelectedNodes[i];
				float result;
				MPI_Status status;
				MPI_Recv(buff, BUFF_SIZE, MPI_PACKED, src, tag, MPI_COMM_WORLD, &status);
				Package pkg;
				pkg.Unpack(buff, BUFF_SIZE);

				//				MPI_Recv(&result, 1, MPI_FLOAT, src, tag, MPI_COMM_WORLD, &status);
				arFunvals[ithPop] = pkg.mData[0];
				//printf("Received response from %dth process: %f\n", src, pkg.mData[0]);
				float* q = pkg.mData;
				vTuple.push_back(NodeTimeTuple(i, q[10]));

				// Write population data
				memset(logStr, 0, 512 * sizeof(char));
				sprintf(logStr, "%d ", iterationCount);
				LogPopData(fplog, logStr);
				for (int j = 0; j < mDim; ++j)
				{
					memset(logStr, 0, 512 * sizeof(char));
					sprintf(logStr, "%f ", pop[ithPop][j]);
					LogPopData(fplog, logStr);
				}
				LogPopData(fplog, "\n");
			}

			sort(vTuple.begin(), vTuple.end());
			for (int i = 0; i < numProcs - 1; ++i)
			{
				memset(logStr, 0, 512 * sizeof(char));
				sprintf(logStr, "Processor number %d took %f seconds.\n", vTuple[i].iNode, vTuple[i].fTime);
				cmaes_WriteLogFile(logStr);
				
// 				int ithPop = ithPass * (numProcs - 1) + vTuple[i].iNode;
// 				for (int j = 0; j < mDim; ++j)
// 				{
// 					memset(logStr, 0, 512 * sizeof(char));
// 					sprintf(logStr, "%f ", pop[ithPop][j]);
// 					cmaes_WriteLogFile(logStr);
// 				}
// 				cmaes_WriteLogFile("\n");
			}
		}

		double currentTime = timer.getCurrentTime();

		/* update the search distribution used for cmaes_SampleDistribution() */
		cmaes_UpdateDistribution(&evo, arFunvals); 

		double* covar = cmaes_GetNew(&evo, "diag(C)");
		double* mean = cmaes_GetNew(&evo, "xmean");
		double* best = cmaes_GetNew(&evo, "xbestever");
		if (covar && mean && best)
		{
			char covarChar[512], meanChar[512], bestChar[512];
			memset(covarChar, 0, 512 * sizeof(char));
			memset(meanChar, 0, 512 * sizeof(char));
			memset(bestChar, 0, 512 * sizeof(char));
			for (int ithDim = 0; ithDim < mDim; ++ithDim)
			{
				sprintf(covarChar, "%s%f  ", covarChar, !info.mUseDegree ? covar[ithDim] : covar[ithDim] / 3.141592653 * 180.0);
				sprintf(meanChar, "%s%f  ", meanChar, !info.mUseDegree ? mean[ithDim] : mean[ithDim] / 3.141592653 * 180.0);
				sprintf(bestChar, "%s%f  ", bestChar, !info.mUseDegree ? best[ithDim] : best[ithDim] / 3.141592653 * 180.0);
			}
			cmaes_WriteLogFile("Diag Covariance:   ");
			cmaes_WriteLogFile(covarChar);
			cmaes_WriteLogFile("\n");
			cmaes_WriteLogFile("Mean:             ");
			cmaes_WriteLogFile(meanChar);
			cmaes_WriteLogFile("\n");
			cmaes_WriteLogFile("Best x so far:    ");
			cmaes_WriteLogFile(bestChar);
			cmaes_WriteLogFile("\n");

			free(covar);
			free(mean);
			free(best);
		}
		iterationCount++;
		memset(logStr, 0, 512 * sizeof(char));
		sprintf(logStr, "-------------------%d CMA iteration finished in %f seconds---------------\n", iterationCount, currentTime);
		cmaes_WriteLogFile(logStr);
		cmaes_WriteLogFile("\n\n");

		if (iterationCount >= maxIterations)
		{

			//			printf("All messages sent.\n");
			break;
		}
		//	/* read instructions for printing output or changing termination conditions */ 
		cmaes_ReadSignals(&evo, info.mCMASignalPath.c_str());   
		fflush(stdout); /* useful in MinGW */
	}
	if (iterationCount < maxIterations)
	{
		const char* terminationText = cmaes_TestForTermination(&evo);
		cmaes_WriteLogFile("Stop:\n");
		cmaes_WriteLogFile(terminationText); /* print termination reason */
	}
	char buff[BUFF_SIZE];
	for (int i = 0; i < numSamples; ++i)
	{
		int dest = pSelectedNodes[i];
		int tag = pSelectedNodes[i];
		Package pkg;
		pkg.Pack(buff, BUFF_SIZE);

		MPI_Send(&buff, BUFF_SIZE, MPI_PACKED, dest, tag, MPI_COMM_WORLD);
	}
	/* get best estimator for the optimum, xmean */
	xfinal = cmaes_GetNew(&evo, "xbestever"); /* "xbestever" might be used as well */
	cmaes_exit(&evo); /* release memory */ 
	memcpy(mPrevSol, xfinal, mDim * sizeof(double));
	memcpy(argMin, xfinal, mDim * sizeof(double));
	/* do something with final solution and finally release memory */
	free(xfinal); 

	fclose(fplog);
}

void CMASearcherMPI::worker(int rank)
{
	printf("Start CMASearchMPI worker function at rank %d\n", rank);
	DecoScene* scene = new DecoScene();
	SceneBuilder::GetSingleton()->ConstructScene(scene);
	printf("Scene constructed.\n");
	while(true)
	{
		MPI_Status status;
		char buff[BUFF_SIZE];
		int src = 0;
		int tag = rank;

		MPI_Recv(buff, BUFF_SIZE, MPI_PACKED, src, tag, MPI_COMM_WORLD, &status);
		Package pkg;
		pkg.Unpack(buff, BUFF_SIZE);
		if (pkg.mNum == 0)
		{
			break;
		}
		//		printf("Message received by process %d is: ", rank);
		//		pkg.Print();
		double* pop = new double[pkg.mNum];
#define RESULT_DETAIL_NUM 11
		double resultDetails[RESULT_DETAIL_NUM];
		for (int i = 0; i < pkg.mNum; ++i)
		{
			pop[i] = pkg.mData[i];
		}
		printf("work %d start evaluation\n", rank);
		DecoTimer timer;
		timer.startTimer();
		float ret = static_cast<float>((*gEvaluator)(pop, scene, resultDetails));
		resultDetails[RESULT_DETAIL_NUM - 1] = timer.getCurrentTime();
		printf("work %d finish evaluation, value: %f\n", rank, ret);
		printf("Params:");
		pkg.Print();
		Package resultPkg(RESULT_DETAIL_NUM, resultDetails);
		resultPkg.Pack(buff, BUFF_SIZE);
		printf("Qs:");
		resultPkg.Print();
		printf("work %d finish packing\n", rank);
		MPI_Send(&buff, BUFF_SIZE, MPI_PACKED, src, tag, MPI_COMM_WORLD);
		printf("work %d finish sending\n", rank);
		delete[] pop;
	}
	printf("worker %d finished his/her work\n", rank);
	delete scene;
}

#endif
